# coding=utf-8

import sys
import logging
from datetime import datetime
from io import BytesIO
from math import log
import platform
import argparse
from typing import Dict, List, Tuple
from pathlib import Path
from threading import Lock

import pandas as pd
from docxtpl import DocxTemplate, InlineImage
from docx.shared import Mm
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
from matplotlib.patches import Patch
import tqdm
from config import (
    RISK_TABLE_PATH, SNPS, RISK_GTS, MALE_ITEMS, FEMALE_ITEMS,
    MALE_TEMPLATE_PATH, FEMALE_TEMPLATE_PATH, RISK_COLORS, LABELS
)

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

plt.style.use('seaborn-v0_8-whitegrid')
SCRIPT_DIR = Path(__file__).parent.resolve()


class FontManager:
    """字体管理类"""
    
    @staticmethod
    def setup_font() -> FontProperties:
        """
        设置中文字体
        Returns:
            FontProperties: 字体属性对象
        """
        try:
            # 首先尝试加载自带字体
            simhei_path = SCRIPT_DIR / 'data' / 'fonts' / 'simhei.ttf'
            if simhei_path.exists():
                return FontProperties(fname=str(simhei_path))
            
            # 在 MacOS 上尝试加载方正兰亭黑
            if platform.system() == 'Darwin':
                fzlth_path = SCRIPT_DIR / 'data' / 'fonts' / '方正兰亭黑.TTF'
                if fzlth_path.exists():
                    return FontProperties(fname=str(fzlth_path))
            
            # 如果没有找到自定义字体，尝试使用系统字体
            return FontProperties(family=['SimHei', 'Microsoft YaHei', 'SimSun'])
            
        except Exception as e:
            logger.warning(f"设置字体时出错: {str(e)}")
            return FontProperties(family=['SimHei', 'Microsoft YaHei', 'SimSun'])


class RiskCalculator:
    """风险计算类"""
    
    def __init__(self):
        self.gt_weights = self._read_risk_lookup_table()
    
    @staticmethod
    def _read_risk_lookup_table() -> Dict:
        """读取风险对照表"""
        try:
            risk_table = pd.read_excel(RISK_TABLE_PATH, header=[0, 1])
            return {
                row[('Unnamed: 0_level_0', '肿瘤名称')]: dict(row)
                for _, row in risk_table.iterrows()
            }
        except Exception as e:
            logger.error(f"读取风险对照表失败: {str(e)}")
            raise

    def get_snp_risk_level(self, sample: Dict) -> Dict[str, str]:
        """计算样本SNP风险等级"""
        return {
            snp: '风险' if ''.join(sorted(sample[snp])) in risk_gts else '正常'
            for snp, risk_gts in zip(SNPS, RISK_GTS)
        }

    def calculate_item_risk(self, sample: Dict) -> Dict[str, Tuple[float, int, float]]:
        """计算样本项目风险
        
        根据样本的SNP基因型计算每个项目的风险分数、风险等级和相对风险比值。
        
        Args:
            sample (Dict): 包含样本SNP基因型信息的字典
            
        Returns:
            Dict[str, Tuple[float, int, float]]: 项目风险信息字典
                - key: 项目名称
                - value: (风险分数, 风险等级, 相对风险比值)的元组
                    - 风险分数: 基于SNP基因型计算的风险得分
                    - 风险等级: 1-4的整数,表示风险等级(1最低,4最高)
                    - 相对风险比值: 相对于人群平均风险的比值,最大为3
        """
        risks = {}
        for item, weights in self.gt_weights.items():
            population_risk = 1.0
            risk_score = 1.0 + sum(
                log(weights[(snp, ''.join(sorted(sample[snp])))])
                for snp in SNPS
            )
            
            risk_level = 1
            if risk_score > population_risk:
                risk_level = 2
            if risk_score > 1.5 * population_risk:
                risk_level = 3
            if risk_score > 2.5 * population_risk:
                risk_level = 4
                
            odds = min(risk_score / population_risk, 3)
            risks[item] = (risk_score, risk_level, odds)
            
        return risks


class ReportGenerator:
    """报告生成类"""
    
    def __init__(self):
        self.font_props = FontManager.setup_font()
        self.risk_calculator = RiskCalculator()
        self.plot_lock = Lock()
        
    def _create_bar_chart(self, items: List[str], risks: Dict) -> BytesIO:
        """创建风险柱状图
        
        根据项目列表和风险信息创建风险等级的柱状图可视化。
        
        Args:
            items (List[str]): 项目名称列表
            risks (Dict): 风险信息字典
                - key: 项目名称
                - value: (风险分数, 风险等级, 相对风险比值)的元组
        
        Returns:
            BytesIO: 包含生成的图表图像数据的字节流对象
            
        Notes:
            - 使用matplotlib创建柱状图
            - y轴范围为0-3,表示相对风险比值
            - 柱状图颜色根据风险等级设置
            - 图例显示四个风险等级(安全/注意/警告/危险)
        """
        with self.plot_lock:
            # 在主线程中创建图表
            plt.switch_backend('Agg')  # 使用非交互式后端
            
            fig, ax = plt.subplots(figsize=(10, 6))
            
            x = items
            y = [risks[item][2] for item in items]
            colors = [RISK_COLORS[risks[item][1] - 1] for item in items]
            
            ax.bar(x, y, color=colors)
            ax.set_ylim(0, 3)
            plt.yticks([0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0])
            
            plt.xticks(rotation=45, ha='right', fontsize='small')
            for label in ax.get_xticklabels():
                label.set_fontproperties(self.font_props)
                
            ax.legend(
                handles=[Patch(facecolor=color) for color in RISK_COLORS],
                labels=LABELS,
                ncol=4, bbox_to_anchor=(0, 1), loc='lower left',
                fontsize='small', prop=self.font_props
            )
            
            image_io = BytesIO()
            fig.savefig(image_io, format='jpg', dpi=250, bbox_inches='tight')
            plt.close(fig)
            return image_io

    def generate_report(self, sample: Dict, output_dir: str) -> None:
        """生成单个样本报告"""
        try:
            # 验证性别
            sample['性别'] = sample['性别'].strip()
            if sample['性别'] not in ['男', '女']:
                raise ValueError(f"无效的性别值: {sample['性别']}")
            
            # 计算风险
            snp_risk_levels = self.risk_calculator.get_snp_risk_level(sample)
            item_risks = self.risk_calculator.calculate_item_risk(sample)
            
            # 选择模板
            items = MALE_ITEMS if sample['性别'] == '男' else FEMALE_ITEMS
            template_path = MALE_TEMPLATE_PATH if sample['性别'] == '男' else FEMALE_TEMPLATE_PATH
            
            # 统计风险等级
            risk_counts = [0] * 4  # [safe, attention, warning, danger]
            for item in items:
                risk_counts[item_risks[item][1] - 1] += 1
            
            # 创建图表
            bar_chart = self._create_bar_chart(items, item_risks)
            
            # 准备上下文数据
            context = {
                'report_id': sample['样本编号'],
                'name': sample['姓名'],
                'sex': sample['性别'],
                'report_date': datetime.today().strftime('%Y{y}%m{m}%d{d}').format(y='年', m='月', d='日'),
                'sample_number': sample['样本编号'],
                **{f"{snp}_gt": sample[snp] for snp in SNPS},
                **{f"{snp}_risk": snp_risk_levels[snp] for snp in SNPS},
                'chief_inspector': sample['主检人'],
                'reviewer': sample['审核人'],
                'danger_count': risk_counts[3],
                'warning_count': risk_counts[2],
                'attention_count': risk_counts[1],
                'save_count': risk_counts[0],
            }
            
            # 渲染报告
            tpl = DocxTemplate(template_path)
            context['risk_bar_image'] = InlineImage(tpl, bar_chart, width=Mm(140), height=Mm(85))
            tpl.render(context=context)
            
            # 保存报告
            output_path = Path(output_dir)
            output_path.mkdir(exist_ok=True)
            report_path = output_path / f"{sample['样本编号']}-{sample['姓名']}.docx"
            tpl.save(str(report_path))
            
            bar_chart.close()
            logger.info(f"成功生成报告: {report_path}")
            
        except Exception as e:
            logger.error(f"生成报告失败 - 样本编号: {sample.get('样本编号', 'unknown')}, 错误: {str(e)}")
            raise


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='生成基因检测报告')
    parser.add_argument('input', help='Excel格式的样本数据文件路径')
    parser.add_argument('output', help='报告输出目录')
    
    args = parser.parse_args()
    
    try:
        input_path = Path(args.input)
        if not input_path.exists():
            raise FileNotFoundError(f"样本文件不存在: {input_path}")
            
        # 读取样本数据
        samples = pd.read_excel(input_path).to_dict('records')
        
        # 生成报告
        report_generator = ReportGenerator()
        for sample in tqdm.tqdm(samples, desc="生成报告"):
            report_generator.generate_report(sample, args.output)
            
        logger.info(f"成功生成 {len(samples)} 份报告")
        
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
        sys.exit(1)


if __name__ == '__main__':
    main()
